Skip to content

Conversation

@siddhesh195
Copy link
Contributor

This patch fixes assertion during tiling due to unsupported linalg.generic containing negative coefficients in indexing expressions by adding a check before attempting tiling

Closes #113021

Fix assertion during tiling due to unsupported linalg.generic containing negative coefficients
@github-actions
Copy link

github-actions bot commented Nov 2, 2024

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir:linalg mlir labels Nov 2, 2024
@llvmbot
Copy link
Member

llvmbot commented Nov 2, 2024

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir-core

Author: Siddhesh Deodhar (siddhesh195)

Changes

This patch fixes assertion during tiling due to unsupported linalg.generic containing negative coefficients in indexing expressions by adding a check before attempting tiling

Closes #113021


Full diff: https://github.com/llvm/llvm-project/pull/114688.diff

4 Files Affected:

  • (modified) mlir/include/mlir/IR/AffineMap.h (+3)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp (+8)
  • (modified) mlir/lib/IR/AffineMap.cpp (+23)
  • (added) mlir/test/Dialect/Linalg/tile-invalid.mlir (+29)
diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index e30950bbf292d6..c2bf0d6b91fed4 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -382,6 +382,9 @@ class AffineMap {
   /// Returns true if the AffineMap represents a symbol-less permutation map.
   bool isPermutation() const;
 
+  /// Returns true if the AffineMap contains non-positive coefficients
+  bool isNonPositiveCoefficients() const;
+
   /// Returns the map consisting of the `resultPos` subset.
   AffineMap getSubMap(ArrayRef<unsigned> resultPos) const;
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index f86715a94b268a..b1e1bf9f721fef 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -119,6 +119,14 @@ struct LinalgOpTilingInterface
     // specified could lead to out of bounds accesses.
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
+    SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+    if (llvm::any_of(linalgOp.getIndexingMapsArray(), [](AffineMap m) {
+          return m.isNonPositiveCoefficients();
+        })) {
+      return linalgOp.emitOpError(
+          "tiling not supported because op has a non positive coefficient");
+    }
+
     SmallVector<Value> valuesToTile = linalgOp->getOperands();
     SmallVector<Value> tiledOperands = makeTiledShapes(
         b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index ea3c0723b07759..63a2506552b28c 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "AffineMapDetail.h"
 #include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -651,6 +652,28 @@ bool AffineMap::isPermutation() const {
   return isProjectedPermutation();
 }
 
+struct CheckCoefficients : public AffineExprVisitor<CheckCoefficients> {
+  CheckCoefficients() {}
+
+  void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
+    visit(expr.getLHS());
+    visit(expr.getRHS());
+    if (expr.getKind() == mlir::AffineExprKind::Mul)
+      isNonPositiveCoefficients |=
+          cast<AffineConstantExpr>(expr.getRHS()).getValue() <= 0;
+  }
+  bool isNonPositiveCoefficients = false;
+};
+
+bool AffineMap::isNonPositiveCoefficients() const {
+
+  return llvm::any_of(getResults(), [](AffineExpr e) {
+    CheckCoefficients t;
+    t.visit(e);
+    return t.isNonPositiveCoefficients;
+  });
+}
+
 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {
   SmallVector<AffineExpr, 4> exprs;
   exprs.reserve(resultPos.size());
diff --git a/mlir/test/Dialect/Linalg/tile-invalid.mlir b/mlir/test/Dialect/Linalg/tile-invalid.mlir
new file mode 100644
index 00000000000000..1684b1de72c264
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/tile-invalid.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics
+
+#map1 = affine_map<(d0) -> (d0)>
+#map2 = affine_map<(d0) -> (-d0 + 7)>
+
+func.func @test(%arg0: tensor<8xi8>, %arg1: tensor<8xi8>) -> tensor<8xi8> {
+   // expected-error @below{{tiling not supported because op has a non positive coefficient}}
+   // expected-error @below{{op faild to tile operation}}
+   // expected-error @below{{op failed to generate tiling loops}}
+  %0 = linalg.generic {
+    indexing_maps = [#map2, #map1],
+    iterator_types = ["parallel"]}
+  ins(%arg0 : tensor<8xi8>)
+  outs(%arg1 : tensor<8xi8>) {
+  ^bb0(%in: i8, %out: i8):
+    linalg.yield %in : i8
+  } -> tensor<8xi8>
+  return %0 : tensor<8xi8>
+}
+
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1:2 = transform.structured.tile_using_for
+      %0 tile_sizes [2] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+     transform.yield
+  }
+}

Copy link
Contributor

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting. I think it "should" work even with negative indices. What is the error you were hitting? (I can try this locally as well in a bit).

@siddhesh195
Copy link
Contributor Author

Interesting. I think it "should" work even with negative indices. What is the error you were hitting? (I can try this locally as well in a bit).

The error we are hitting is due to an assertion here

@MaheshRavishankar
Copy link
Contributor

Interesting. I think it "should" work even with negative indices. What is the error you were hitting? (I can try this locally as well in a bit).

The error we are hitting is due to an assertion here

What would happen if you remove that assertion? I just wanted to see what the IR looks like if you just let it tile this way?

@christopherbate
Copy link
Contributor

christopherbate commented Nov 14, 2024

Removing the assertion will not result in a correct result. Here is the result from test case:

within split at test.mlir:1 offset :8:8: error: 'linalg.generic' op inferred input/output operand #0 has shape's dimension #0 to be greater than or equal to 8, but found 7
  %0 = linalg.generic {
       ^
within split at test.mlir:1 offset :8:8: note: see current operation: 
%7 = "linalg.generic"(%5, %6) <{indexing_maps = [affine_map<(d0) -> (-d0 + 7)>, affine_map<(d0) -> (d0)>], iterator_types = [#linalg.iterator_type<parallel>], operandSegmentSizes = array<i32: 1, 1>}> ({
^bb0(%arg4: i8, %arg5: i8):
  "linalg.yield"(%arg4) : (i8) -> ()
}) : (tensor<7xi8>, tensor<2xi8>) -> tensor<2xi8>
// -----// IR Dump After InterpreterPass Failed (transform-interpreter) //----- //
#map = affine_map<(d0) -> (-d0 + 7)>
#map1 = affine_map<(d0) -> (d0)>
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<8xi8>, tensor<8xi8>) -> tensor<8xi8>, sym_name = "test"}> ({
  ^bb0(%arg1: tensor<8xi8>, %arg2: tensor<8xi8>):
    %2 = "arith.constant"() <{value = 0 : index}> : () -> index
    %3 = "arith.constant"() <{value = 8 : index}> : () -> index
    %4 = "arith.constant"() <{value = 2 : index}> : () -> index
    %5 = "scf.for"(%2, %3, %4, %arg2) ({
    ^bb0(%arg3: index, %arg4: tensor<8xi8>):
      %6 = "affine.apply"(%arg3) <{map = #map}> : (index) -> index
      %7 = "tensor.extract_slice"(%arg1, %6) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>, static_offsets = array<i64: -9223372036854775808>, static_sizes = array<i64: 7>, static_strides = array<i64: 1>}> : (tensor<8xi8>, index) -> tensor<7xi8>
      %8 = "tensor.extract_slice"(%arg4, %arg3) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>, static_offsets = array<i64: -9223372036854775808>, static_sizes = array<i64: 2>, static_strides = array<i64: 1>}> : (tensor<8xi8>, index) -> tensor<2xi8>
      %9 = "linalg.generic"(%7, %8) <{indexing_maps = [#map, #map1], iterator_types = [#linalg.iterator_type<parallel>], operandSegmentSizes = array<i32: 1, 1>}> ({
      ^bb0(%arg5: i8, %arg6: i8):
        "linalg.yield"(%arg5) : (i8) -> ()
      }) : (tensor<7xi8>, tensor<2xi8>) -> tensor<2xi8>
      %10 = "tensor.insert_slice"(%9, %arg4, %arg3) <{operandSegmentSizes = array<i32: 1, 1, 1, 0, 0>, static_offsets = array<i64: -9223372036854775808>, static_sizes = array<i64: 2>, static_strides = array<i64: 1>}> : (tensor<2xi8>, tensor<8xi8>, index) -> tensor<8xi8>
      "scf.yield"(%10) : (tensor<8xi8>) -> ()
    }) : (index, index, index, tensor<8xi8>) -> tensor<8xi8>
    "func.return"(%5) : (tensor<8xi8>) -> ()
  }) : () -> ()
  "builtin.module"() ({
    "transform.named_sequence"() <{arg_attrs = [{transform.readonly}], function_type = (!transform.any_op) -> (), sym_name = "__transform_main"}> ({
    ^bb0(%arg0: !transform.any_op):
      %0 = "transform.structured.match"(%arg0) <{ops = ["linalg.generic"]}> : (!transform.any_op) -> !transform.any_op
      %1:2 = "transform.structured.tile_using_for"(%0) <{scalable_sizes = array<i1: false>, static_sizes = array<i64: 2>}> : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
      "transform.yield"() : () -> ()
    }) : () -> ()
  }) {transform.with_named_sequence} : () -> ()
}) : () -> ()

Copy link
Contributor

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok. Thanks for the due diligence.

@siddhesh195
Copy link
Contributor Author

Ok. Thanks for the due diligence.

Thank you for the review! If the patch looks good, can you help merge it since I do not have write access

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mlir:core MLIR Core Infrastructure mlir:linalg mlir

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[mlir][linalg] Tiling implementation for linalg.generic will result in assert/crash if indexing expressions have negative multiplicative coefficients

4 participants